from typing import Dict, Iterable, List, Optional

import torch
from torch import nn


class NNLoss(nn.Module):
    """
    兼容多任务模型需求的损失函数基类
    """

    def __init__(
        self,
        sub_loss_names: List[str],
        sub_loss_weights: Optional[List[float]] = None,
    ):
        """
        :param sub_loss_names:
            子损失的名称
        :param sub_loss_weights:
            子损失权重
        """
        super().__init__()
        self.names = sub_loss_names
        self.weights = (
            sub_loss_weights if sub_loss_weights is not None else [1] * len(sub_loss_names)
        )
        assert len(self.names) == len(self.weights)

    def forward(self, outputs, targets=None) -> Dict[str, torch.Tensor]:
        # 推理产生loss，将结果封装成字典，并运用子损失权重进行加权从而产生总损失
        loss_dict = dict()
        losses = self.calculate_loss(outputs, targets)
        total_loss = torch.tensor(0.0, requires_grad=True)
        for loss, name, weight in zip(losses, self.names, self.weights):
            loss_dict[name] = loss
            if not (
                loss is None
                or torch.isinf(loss)
                or torch.isnan(loss)
                or (loss == 0 and loss.grad_fn is None)
            ):
                if total_loss.device != loss.device:
                    total_loss = total_loss.to(loss.device)
                total_loss = total_loss + loss * weight
        loss_dict["total_loss"] = total_loss
        return loss_dict

    def calculate_loss(self, outputs, targets) -> Iterable[torch.Tensor]:
        """
        :param outputs:
            模型推理输出的张量
        :param targets:
            模型预测目标（由NNDataset的build_target构建而成）
        :return:
            与self.names顺序和数量相同的子损失序列
        """
        raise NotImplementedError
